# memory_system/memory_manager.py
"""
增强版记忆管理器 - 修复版本
支持多种记忆类型和优化检索算法
"""
import chromadb
from datetime import datetime, timedelta
import json
import os
import sys

current_dir = os.path.dirname(__file__)
if current_dir not in sys.path:
    sys.path.append(current_dir)

from chroma_setup import setup_chroma_database

class MemoryManager:
    """增强版记忆管理器 - 支持多种记忆类型和优化检索"""
    
    # 定义记忆类型常量
    MEMORY_TYPES = {
        "interaction": "用户交互记录",
        "mistake": "错题集", 
        "progress": "学习进度",
        "derived_knowledge": "衍生知识",
        "user_profile": "用户特征"
    }
    
    def __init__(self, db_path="./memory_db"):
        print("🚀 初始化增强版记忆管理器...")
        try:
            self.client, self.collection = setup_chroma_database(db_path)
            if not self.collection:
                raise Exception("数据库连接失败")
            print("✅ 记忆管理器初始化成功！")
        except Exception as e:
            print(f"❌ 记忆管理器初始化失败: {e}")
            raise
    
    def write_memory(self, user_id, content, memory_type="interaction", 
                    metadata=None, importance=0.5):
        """写入记忆到数据库 - 支持多种记忆类型"""
        if metadata is None:
            metadata = {}
        
        try:
            # 生成唯一记忆ID
            memory_id = f"{user_id}_{memory_type}_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
            
            # 准备元数据 - 确保没有 None 值
            full_metadata = {
                "user_id": user_id,
                "memory_type": memory_type,
                "timestamp": datetime.now().isoformat(),
                "importance": str(importance),
            }
            
            # 处理传入的元数据 - 过滤掉 None 值并转换为字符串
            for key, value in metadata.items():
                if value is None:
                    continue  # 跳过 None 值
                elif isinstance(value, list):
                    full_metadata[key] = ",".join(str(v) for v in value)
                elif isinstance(value, (str, int, float, bool)):
                    full_metadata[key] = str(value)
                else:
                    full_metadata[key] = str(value)
            
            # 写入数据库
            self.collection.add(
                documents=[content],
                metadatas=[full_metadata],
                ids=[memory_id]
            )
            
            print(f"✅ [{memory_type}] 记忆已保存 - 用户: {user_id}")
            return {
                "status": "success", 
                "memory_id": memory_id,
                "memory_type": memory_type,
                "message": f"{self.MEMORY_TYPES.get(memory_type, '未知类型')}保存成功"
            }
            
        except Exception as e:
            error_msg = f"写入记忆失败: {e}"
            print(f"❌ {error_msg}")
            return {"status": "error", "message": error_msg}
    
    def write_mistake(self, user_id, question, wrong_answer, 
                     correct_answer, knowledge_points, 
                     difficulty="medium", analysis=""):
        """写入错题记录"""
        content = f"错题: {question} | 错误答案: {wrong_answer} | 正确答案: {correct_answer}"
        if analysis:
            content += f" | 分析: {analysis}"
            
        metadata = {
            "question": question,
            "wrong_answer": wrong_answer,
            "correct_answer": correct_answer,
            "knowledge_points": knowledge_points,
            "difficulty": difficulty,
            "analysis": analysis,
            "mistake_count": "1",
            "last_mistake_time": datetime.now().isoformat()
        }
        
        return self.write_memory(
            user_id=user_id,
            content=content,
            memory_type="mistake",
            metadata=metadata,
            importance=0.8
        )
    
    def write_progress(self, user_id, course, chapter, 
                      progress, confidence, notes=""):
        """写入学习进度"""
        content = f"课程: {course} | 章节: {chapter} | 进度: {progress*100}% | 掌握度: {confidence*100}%"
        if notes:
            content += f" | 笔记: {notes}"
            
        metadata = {
            "course": course,
            "chapter": chapter,
            "progress": str(progress),
            "confidence": str(confidence),
            "notes": notes,
        }
        
        # 只有完成时才添加完成日期
        if progress >= 1.0:
            metadata["completion_date"] = datetime.now().isoformat()
        
        return self.write_memory(
            user_id=user_id,
            content=content,
            memory_type="progress",
            metadata=metadata,
            importance=0.6
        )
    
    def read_memory(self, user_id, memory_type=None, limit=10):
        """读取用户记忆 - 简化版本，避免复杂查询"""
        try:
            # 获取所有该用户的记忆
            results = self.collection.get(
                where={"user_id": user_id},
                limit=limit * 3  # 获取更多然后过滤
            )
            
            memories = []
            if results['ids']:
                for i in range(len(results['ids'])):
                    memory_data = results['metadatas'][i]
                    # 在内存中过滤记忆类型
                    if memory_type is None or memory_data.get("memory_type") == memory_type:
                        memory = {
                            "id": results['ids'][i],
                            "content": results['documents'][i],
                            "metadata": memory_data,
                            "type": memory_data.get("memory_type", "unknown")
                        }
                        memories.append(memory)
                
                # 按时间排序，最新的在前面
                memories.sort(key=lambda x: x["metadata"].get("timestamp", ""), reverse=True)
                # 限制返回数量
                memories = memories[:limit]
            
            type_desc = f"{memory_type}类型" if memory_type else ""
            print(f"✅ 为用户 {user_id} 找到 {len(memories)} 条{type_desc}记忆")
            return memories
            
        except Exception as e:
            error_msg = f"读取记忆失败: {e}"
            print(f"❌ {error_msg}")
            return []
    
    def search_similar_memories(self, query_text, user_id=None, 
                               memory_type=None, limit=5, 
                               similarity_threshold=0.7):
        """基于向量相似性搜索记忆"""
        try:
            # 简单的相似度搜索，不过滤类型
            results = self.collection.query(
                query_texts=[query_text],
                n_results=limit * 2
            )
            
            similar_memories = []
            if results['ids'] and results['ids'][0]:
                for i in range(len(results['ids'][0])):
                    # 计算相似度分数
                    distance = results['distances'][0][i] if results['distances'] else 0
                    similarity_score = 1.0 / (1.0 + distance)
                    
                    memory_data = results['metadatas'][0][i]
                    
                    # 应用过滤条件
                    if user_id and memory_data.get("user_id") != user_id:
                        continue
                    if memory_type and memory_data.get("memory_type") != memory_type:
                        continue
                    if similarity_score < similarity_threshold:
                        continue
                    
                    memory = {
                        "id": results['ids'][0][i],
                        "content": results['documents'][0][i],
                        "metadata": memory_data,
                        "similarity_score": round(similarity_score, 3),
                        "type": memory_data.get("memory_type", "unknown")
                    }
                    similar_memories.append(memory)
                
                # 按相似度排序并限制数量
                similar_memories.sort(key=lambda x: x["similarity_score"], reverse=True)
                similar_memories = similar_memories[:limit]
            
            print(f"✅ 找到 {len(similar_memories)} 条相似记忆")
            return similar_memories
            
        except Exception as e:
            print(f"❌ 相似记忆搜索失败: {e}")
            return []
    
    def get_user_mistakes(self, user_id, knowledge_point=None, recent_days=30):
        """获取用户错题集 - 简化版本"""
        try:
            # 获取所有错题
            mistakes = self.read_memory(user_id, memory_type="mistake", limit=50)
            
            # 按时间过滤
            since_date = datetime.now() - timedelta(days=recent_days)
            recent_mistakes = []
            for mistake in mistakes:
                timestamp_str = mistake["metadata"].get("timestamp", "")
                if timestamp_str:
                    try:
                        timestamp = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
                        if timestamp >= since_date:
                            recent_mistakes.append(mistake)
                    except:
                        recent_mistakes.append(mistake)  # 如果解析失败，包含该记录
            
            # 知识点过滤
            if knowledge_point:
                filtered_mistakes = []
                for mistake in recent_mistakes:
                    k_points = mistake["metadata"].get("knowledge_points", "")
                    if knowledge_point in k_points:
                        filtered_mistakes.append(mistake)
                recent_mistakes = filtered_mistakes
            
            # 按错误次数和时间排序
            recent_mistakes.sort(key=lambda x: (
                -int(x["metadata"].get("mistake_count", 1)),
                x["metadata"].get("timestamp", "")
            ), reverse=True)
            
            print(f"✅ 找到 {len(recent_mistakes)} 条错题记录")
            return recent_mistakes
            
        except Exception as e:
            print(f"❌ 获取错题集失败: {e}")
            return []
    
    def get_learning_progress(self, user_id, course=None):
        """获取用户学习进度统计"""
        try:
            # 获取所有进度记录
            progress_records = self.read_memory(user_id, memory_type="progress", limit=100)
            
            if not progress_records:
                return {"overall_progress": 0, "course_progress": {}}
            
            # 分析各课程进度
            course_progress = {}
            for record in progress_records:
                metadata = record["metadata"]
                course_name = metadata.get("course", "未知课程")
                chapter = metadata.get("chapter", "未知章节")
                
                # 安全地转换进度和置信度
                try:
                    progress = float(metadata.get("progress", 0))
                except:
                    progress = 0
                
                try:
                    confidence = float(metadata.get("confidence", 0))
                except:
                    confidence = 0
                
                if course_name not in course_progress:
                    course_progress[course_name] = {
                        "chapters": [],
                        "average_progress": 0,
                        "average_confidence": 0
                    }
                
                course_progress[course_name]["chapters"].append({
                    "chapter": chapter,
                    "progress": progress,
                    "confidence": confidence,
                    "timestamp": metadata.get("timestamp")
                })
            
            # 计算平均进度和掌握度
            overall_progress = 0
            course_count = 0
            
            for course_name, data in course_progress.items():
                chapters = data["chapters"]
                if chapters:
                    avg_progress = sum(ch["progress"] for ch in chapters) / len(chapters)
                    avg_confidence = sum(ch["confidence"] for ch in chapters) / len(chapters)
                    
                    course_progress[course_name]["average_progress"] = round(avg_progress, 3)
                    course_progress[course_name]["average_confidence"] = round(avg_confidence, 3)
                    
                    overall_progress += avg_progress
                    course_count += 1
            
            overall_progress = round(overall_progress / course_count, 3) if course_count > 0 else 0
            
            return {
                "overall_progress": overall_progress,
                "course_progress": course_progress,
                "total_courses": course_count,
                "analysis_date": datetime.now().isoformat()
            }
            
        except Exception as e:
            print(f"❌ 获取学习进度失败: {e}")
            return {}
    
    def get_knowledge_gap_analysis(self, user_id):
        """知识缺口分析 - 基于错题和学习进度"""
        try:
            # 获取近期错题
            recent_mistakes = self.get_user_mistakes(user_id, recent_days=30)
            
            # 分析错误知识点
            knowledge_weaknesses = {}
            for mistake in recent_mistakes:
                metadata = mistake["metadata"]
                k_points = metadata.get("knowledge_points", "")
                if k_points:
                    points = k_points.split(",")
                    for point in points:
                        point = point.strip()
                        if point:
                            knowledge_weaknesses[point] = knowledge_weaknesses.get(point, 0) + 1
            
            # 获取学习进度
            progress_data = self.get_learning_progress(user_id)
            
            # 生成知识缺口报告
            weak_points = sorted(knowledge_weaknesses.items(), key=lambda x: x[1], reverse=True)[:5]
            
            analysis = {
                "user_id": user_id,
                "analysis_date": datetime.now().isoformat(),
                "weak_knowledge_points": [
                    {"point": point, "mistake_count": count} 
                    for point, count in weak_points
                ],
                "learning_progress": progress_data,
                "recommendations": self._generate_study_recommendations(weak_points, progress_data)
            }
            
            print(f"✅ 完成知识缺口分析，发现 {len(weak_points)} 个薄弱知识点")
            return analysis
            
        except Exception as e:
            print(f"❌ 知识缺口分析失败: {e}")
            return {}
    
    def _generate_study_recommendations(self, weak_points, progress_data):
        """生成学习建议"""
        recommendations = []
        
        if weak_points:
            top_weak = weak_points[0][0]
            recommendations.append(f"重点复习知识点: {top_weak}")
            
            if len(weak_points) > 1:
                other_weak = ", ".join([p[0] for p in weak_points[1:3]])
                recommendations.append(f"同时关注: {other_weak}")
        
        # 基于进度给出建议
        overall_progress = progress_data.get("overall_progress", 0)
        if overall_progress < 0.3:
            recommendations.append("学习进度较慢，建议增加学习频率")
        elif overall_progress > 0.8:
            recommendations.append("学习进度良好，建议进行综合复习")
        
        return recommendations

# 测试增强功能
def test_enhanced_memory_system():
    """测试增强版记忆系统"""
    print("🧪 测试增强版记忆系统...")
    
    try:
        # 创建记忆管理器
        mm = MemoryManager()
        
        # 测试错题记录
        print("\n📝 测试错题记录...")
        mistake_result = mm.write_mistake(
            user_id="med_student_001",
            question="卷积神经网络中池化层的主要作用是什么？",
            wrong_answer="增加网络深度",
            correct_answer="降维和特征不变性",
            knowledge_points=["CNN", "池化层", "特征提取"],
            difficulty="medium",
            analysis="混淆了池化层和卷积层的作用"
        )
        print(f"错题记录结果: {mistake_result}")
        
        # 测试学习进度 - 使用安全的参数
        print("\n📊 测试学习进度记录...")
        progress_result = mm.write_progress(
            user_id="med_student_001",
            course="医学人工智能",
            chapter="深度学习基础",
            progress=0.75,
            confidence=0.8,
            notes="掌握了CNN基本原理，还需练习调参"
        )
        print(f"进度记录结果: {progress_result}")
        
        # 测试基础记忆检索
        print("\n🔍 测试基础记忆检索...")
        all_memories = mm.read_memory("med_student_001", limit=10)
        print(f"找到 {len(all_memories)} 条总记忆")
        for i, memory in enumerate(all_memories, 1):
            print(f"  {i}. [{memory['type']}] {memory['content'][:30]}...")
        
        # 测试错题检索
        print("\n📝 测试错题检索...")
        mistakes = mm.get_user_mistakes("med_student_001")
        print(f"找到 {len(mistakes)} 条错题记录")
        
        # 测试进度分析
        print("\n📈 测试学习进度分析...")
        progress_data = mm.get_learning_progress("med_student_001")
        print(f"整体进度: {progress_data.get('overall_progress', 0) * 100:.1f}%")
        
        # 测试知识缺口分析
        print("\n🎯 测试知识缺口分析...")
        gap_analysis = mm.get_knowledge_gap_analysis("med_student_001")
        print("知识缺口分析完成!")
        if gap_analysis.get("weak_knowledge_points"):
            print("薄弱知识点:")
            for weak in gap_analysis["weak_knowledge_points"]:
                print(f"  - {weak['point']} (错误次数: {weak['mistake_count']})")
        
        print("\n🎉 记忆系统测试完成！")
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_enhanced_memory_system()